
from langchain_core.output_parsers import StrOutputParser
from langchain_openai import ChatOpenAI
from langchain_core.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough,RunnableLambda
from langchain_huggingface import HuggingFaceEmbeddings

# 创建llm模型
def init_model() :
    # 模型定义
    ZHIPUAI_API_KEY = '794380a4cee054a0f96bb2844b41fd12.X4t70kph1CfmoKfT'
    BASE_PATH ='https://open.bigmodel.cn/api/paas/v4/'
    return ChatOpenAI(model_name='glm-4',temperature=.7,openai_api_key=ZHIPUAI_API_KEY,base_url=BASE_PATH)

# 文档向量化
def init_embeddings():
    # pip install   sentence_transformers 
    # from langchain_community.embeddings import HuggingFaceEmbeddings
    import os
    os.environ['HF_ENDPOINT'] = 'https://hf-mirror.com'
    return HuggingFaceEmbeddings(model_name="google-bert/bert-base-chinese")

# 模板路由
def prompt_route(input):
    py_tmpl = """你是一个非常优秀的物理学家，你擅长回答物理问题，你之所以那么优秀是因为你能将难题分解成不同的组成部分，回答组成部分的问题，然后把他们结合起来回答更广泛的问题， 这是一个问题{query}"""
    ma_tmpl = """你是一个非常优秀的数学家，你擅长回答数学问题，你之所以那么优秀是因为你能将难题分解成不同的组成部分，回答组成部分的问题，然后把他们结合起来回答更广泛的问题， 这是一个问题{query}"""
    prompt_templates = [py_tmpl,ma_tmpl]
    embeddings = init_embeddings()
    prompt_embeddings = embeddings.embed_documents(prompt_templates)
    query_embedding = embeddings.embed_query(input["query"])
    from langchain.utils.math import  cosine_similarity
    similarity = cosine_similarity([query_embedding],prompt_embeddings)[0]
    print('max',similarity.argmax())
    most_similar = prompt_templates[similarity.argmax()]   
    print("使用数学" if most_similar == ma_tmpl else "使用物理")
    return PromptTemplate.from_template(most_similar)

chain = (
    {"query":RunnablePassthrough()}
    | RunnableLambda(prompt_route)
    | init_model() 
    | StrOutputParser()
)

chain.invoke('什么是黑洞')
# 执行结果：
# max 0
# 使用物理
chain.invoke('什么是勾股定理')
# 执行结果：
# max 1
# 使用数学
